#include "FreeImage.h"

#include <iostream>

char *g_imageName;
int g_regionSize;
FIBITMAP *g_src;
FIBITMAP *g_dst;
int g_width;
int g_height;
int g_bpp;
int g_nRegions;

bool g_isUsed[100]; //indicates whether a given strip index is used or not
int g_left = 0; //left-most strip index
int g_right = 0; //right-most strip index
int g_result[100]; //stores the strip indices 

#define ABS(x) x < 0 ? -x : x

#define N_HORIZONTAL 4		//how many pixels on either end of the "strip boundary" to average together
unsigned int ScoreRegions(int lRegion, int rRegion, int stripSize = g_regionSize)
{
	int bpp = g_bpp/8;
	unsigned int score = 0;
	for(int y=0; y<g_height; y++) {
		int sumLeft = 0;
		int sumRight = 0;
		for(int x=0; x<N_HORIZONTAL; x++) {
			BYTE *p1 = FreeImage_GetScanLine(g_src, y) + (stripSize*lRegion+stripSize-1-x)*bpp;
			BYTE *p2 = FreeImage_GetScanLine(g_src, y) + (stripSize*rRegion+x)*bpp;
			sumLeft += p1[FI_RGBA_RED];
			sumLeft += p1[FI_RGBA_BLUE];
			sumLeft += p1[FI_RGBA_GREEN];			
			sumRight += p2[FI_RGBA_RED];
			sumRight += p2[FI_RGBA_BLUE];
			sumRight += p2[FI_RGBA_GREEN];
		}
		sumLeft /= 3*N_HORIZONTAL;	//normalize by number of color channels and number of pixels.  shouldn't be necessary, but just to be safe
		sumRight /= 3*N_HORIZONTAL;
		score += (sumLeft-sumRight)*(sumLeft-sumRight);
	}
	return score;
}

void BlankDst()
{
	for(int y=0; y<g_height; y++) {
		BYTE *pDst = FreeImage_GetScanLine(g_dst, y);
		memset(pDst, 0, g_width*g_bpp/8);
	}
}

void PasteRegionRange(int dstStart, int srcStart, int srcEnd)
{
	int bpp = g_bpp/8;
	int pasteWidth = (srcEnd-srcStart)*g_regionSize;
	for(int y = 0; y < g_height; y++) {
		BYTE *pSrc = FreeImage_GetScanLine(g_src, y) + srcStart*bpp*g_regionSize;
		BYTE *pDst = FreeImage_GetScanLine(g_dst, y) + dstStart*bpp*g_regionSize;

		memcpy(pDst, pSrc, pasteWidth*bpp);
	}
}

int FindStripSize()
{
	unsigned int worstScore = 1;
	int bestGuess = 1;
	//margin of error = lowest ratio between current worst score and the score of all other strip sizes smaller than current best guess 
	//that aren't multiples of current best guess.  margin of error is
	//used for ruling out strip sizes that are just multiples of the "real" strip size.
	float marginOfError = 100000.0f;		
	for(int stripSize=4; stripSize<g_width/2; stripSize++) {
		if(g_width%stripSize != 0)		//stripSize must be a factor of image width
			continue;
		unsigned int score = 0;
		int nStrips = g_width/stripSize;
		for(int strip=0; strip<nStrips-1; strip++) {
			score += ScoreRegions(strip, strip+1, stripSize);
		}
		score = score/nStrips;
		if(score > worstScore) {
			bool justMultiple = false;
			if(stripSize % bestGuess == 0) {
				if((float)score/worstScore < marginOfError) 
					justMultiple = true;		//this is just a multiple of the "real" strip size so don't count it as buest guess
			} 
			if(!justMultiple) {
				worstScore = score;
				bestGuess = stripSize;
			}
		}
		if(stripSize%bestGuess != 0) {
			float moe = (float)worstScore/score;
			if(moe < marginOfError)		//update margin of error as lowest ratio of score between current best guess and any other score
				marginOfError = moe;
		}
	}
	return bestGuess;
}

int main(int argc, char *argv[])
{
	if(argc <2) {
		std::cout << "Usage:  Instagram_Challenge.exe image_path" << std::endl;
		return 0;
	}
	g_imageName = argv[1];

	g_src = FreeImage_Load(FIF_PNG, g_imageName);
	if(!g_src) {
		std::cout << "Unable to load supplied file." << std::endl;
		return -1;
	}

	g_width = FreeImage_GetWidth(g_src);
	g_height = FreeImage_GetHeight(g_src);
	g_bpp = FreeImage_GetBPP(g_src);

	g_dst = FreeImage_Allocate(g_width, g_height, g_bpp);

	std::cout << "Calculating strip size... " << std::flush;
	g_regionSize = FindStripSize();		//BONUS!!
	g_nRegions = g_width/g_regionSize;
	std::cout << g_regionSize << std::endl;


	//Initialize Result
	for(int i=0; i<sizeof(g_result)/sizeof(g_result[0]); i++)
		g_result[i] = i;
	for(int i=0; i<sizeof(g_isUsed)/sizeof(g_isUsed[0]); i++)
		g_isUsed[i] = false;

	//lock in region 0 at place 0 for now
	g_left = g_right = 0;		//grow out out from 0 from both left and right
	g_isUsed[0] = true;
	g_result[g_left] = 0;

	std::cout << "Unshredding... " << std::flush;
	//grow our unshredded image result on both left and right based on what's the best "fitting" strip
	//on the left or right of our running result.  each iteration adds a strip to our result
	for(int region=1; region<g_nRegions; region++) {
		unsigned int bestScore = 4294967295;			//Initialize to highest possible value
		int bestRgn = -1;
		bool isLeft = true;
		for(int rgn=0; rgn<g_nRegions; rgn++) {
			if(g_isUsed[rgn])		
				continue;//already part of result.  skip it
			unsigned int score = ScoreRegions(rgn, g_result[g_left]);
			if(score < bestScore) {
				isLeft = true;
				bestScore = score;
				bestRgn = rgn;
			}
			score = ScoreRegions(g_result[g_right], rgn);
			if(score < bestScore) {
				isLeft = false;
				bestScore = score;
				bestRgn = rgn;
			}
		}
		int idx;
		if(isLeft) {
			g_left--;	 
			if(g_left < 0)
				g_left += g_nRegions;
			idx = g_left;
		} else {
			g_right++;
			if(g_right >= g_nRegions)
				g_right -= g_nRegions;
			idx = g_right;
		}
		g_isUsed[bestRgn] = true;
		g_result[idx] = bestRgn;
	}

	for(int rgn=0; rgn<g_nRegions; rgn++) {
		int idx = g_left+rgn;
		if(idx >= g_nRegions)
			idx -= g_nRegions;
		PasteRegionRange(rgn, g_result[idx], g_result[idx]+1);
	}
	std::cout << "Done." << std::endl;

	FreeImage_Save(FIF_BMP, g_dst, "result.bmp");
	std::cout << "Result saved as 'result.bmp'" << std::endl;

	FreeImage_Unload(g_dst);
	FreeImage_Unload(g_src);

	return 0;
}